import numpy as np
from sklearn import datasets
from sklearn.model_selection import train_test_split

# 加载鸢尾花数据集
iris = datasets.load_iris()
X = iris.data
Y = iris.target

# 分割训练集和测试集，random_state=4保证每次运行结果一致
x_train, x_test, y_train, y_test = train_test_split(X, Y, random_state=4)
print('训练集样本:',(x_train)) # 打印训练集样本
print('测试集样本:',(x_test))  # 打印测试集样本
print('训练集标签:',(y_train)) # 打印训练集标签
print('测试集标签:',(y_test))  # 打印测试集标签

# 计算每一类的先验概率，以及每一类中每一种特征的均值和方差
def prior_Mean_Var(x_train, y_train):
    prior_list = []  # 存储每个类别的先验概率
    mean = np.array([[0, 0, 0, 0]])  # 初始化均值数组
    var = np.array([[0, 0, 0, 0]])   # 初始化方差数组
    for kind in range(3):  # 遍历三个类别
        x_class = x_train[np.nonzero(kind == y_train)]  # 获取当前类别的所有样本
        prior_list.append(len(x_class) / len(x_train))  # 计算并添加先验概率
        m = np.mean(x_class, axis=0, keepdims=True)  # 计算当前类别的均值
        mean = np.append(mean, m, axis=0)  # 将均值添加到均值数组中
        v = np.var(x_class, axis=0, keepdims=True)  # 计算当前类别的方差
        var = np.append(var, v, axis=0)  # 将方差添加到方差数组中
    return prior_list, mean[1:], var[1:]  # 返回先验概率列表、均值数组和方差数组（去掉初始化的第一行）

# 对样本进行分类
def predict(x_test, y_test, prior, mean, var):
    _class = []  # 存储预测的类别
    eps = 1e-10  # 防止分母为0的小常数
    for i in x_test:
        x = np.tile(i, (3, 1))  # 将测试样本复制三份，对应三个类别
        p = (np.exp(-(x - mean) ** 2 / (2 * var + eps))) / (np.sqrt(2 * np.pi) * var + eps)  # 计算高斯分布的概率密度函数值
        p_after = np.sum(np.log(p), axis=1)  # 计算后验概率的对数值
        p_class = np.log(prior) + p_after  # 计算每个类别的最终概率
        _class.append(np.argmax(p_class))  # 选择概率最大的类别作为预测结果
    return _class

# 计算先验概率、均值和方差
prior, mean, var = prior_Mean_Var(x_train, y_train)
# 使用训练好的模型对测试集进行预测
y_pred = predict(x_test, y_test, prior, mean, var)

# 计算模型的准确率：
count = 0
for i in range(len(y_test)):
    if y_pred[i] == y_test[i]:
        count += 1
accuracy = count / len(y_test)  # 计算准确率
print('accuracy: {:.2%}'.format(accuracy))  # 打印准确率，格式化为百分比形式
